package immutablecollections;

import immutablecollections.exceptions.*;
import immutablecollections.functions.*;
import immutablecollections.misc.*;
import immutablecollections.textboxes.*;
import java.util.*;

/**
 * An immutable version of a Rose Tree (Multi-way Tree or just Tree) - a tree in which each node has a *value* and
 * an arbitrary number of *sub-trees*.
 * 
 * Note that RoseTrees are not the same as B-Trees.
 * 
 * ## Introduction
 * 
 * <img src="doc-files/classes-focus-on-rosetree.png" alt="classes-focus-on-rosetree" style="width: 40em"/></p>
 * 
 * A rose tree has a root node containing an element and an ordered list (shelf) of sub-trees.
 * The definition is recursive - each sub-tree
 * also has an element and a shelf of
 * sub-trees and so on. Leaf trees are those with no sub-trees - ie the shelf is empty.
 * 
 * For example, the tree with element `a` and three sub nodes:
 * 
 * 1. The leaf node `b`
 * 2. The tree with element `c` and three leaf nodes `e`, `f`,`g`
 * 3. The tree with element `d` and the sub tree of the leaf node `h`
 * 
 * <img src="doc-files/rose-tree-a.png" alt="rose-tree-a"  style="width: 20%"/></p>
 * 
 * We could also represent it like this:
 * 
 * <img src="doc-files/rose-tree-b.png" alt="rose-tree-b"  style="width: 20%"/></p>
 * 
 * or like this:
 * 
 *     a        
 *     .........
 *     b c     d
 *       ..... .
 *       e f g h
 * 
 * This is an example of the what {@link #toBoxString()} produces.
 * 
 * ## Creation
 * 
 * - {@link #withElements(Object, Object...)}
 * - {@link #withElements(Object, Collection)}
 * - {@link #withNodes(Object, ImRoseTree...)}
 * - {@link #withNodes(Object, Collection)}
 * 
 * ## Query
 * 
 * - {@link #contains(Object)}
 * - {@link #getElement()}
 * - {@link #getSubTrees()}
 * - {@link #getNodeAtIndex(int)}
 * - {@link #size()}
 * 
 * ## Mutation
 * 
 * - {@link #replaceElement(Object)}
 * - {@link #map(Function1)}
 * 
 * ## Iteration
 * 
 * - {@link #iterator()}
 * - {@link #getZipper()}
 * - {@link #getZipperIterator()}
 * 
 * ## String representation
 * 
 * - {@link #toString()}
 * - {@link #toBoxString()}
 * 
 * ## Implementation
 * 
 * The usual functional implementation of a rose tree would be to represent the children of a node by a (functional
 * style) list. We have chosen a slightly different approach that uses an `ImShelf` instead. This gives better
 * performance when mutations to the children are randomly distributed.
 * 
 * ## Notation
 * 
 * In this library our convention is to call a *binary* tree a *Tree* and a *multi-way* tree a *Rose Tree*.
 * 
 * In the examples below we are assuming that `ImRoseTree` has been statically imported:
 * 
 *     import static immutablecollections.ImRoseTree.*;
 *     
 * which means that we can use (eg)
 * 
 *     withNodes(1, leaf(2), leaf(3))
 * 
 * instead of
 * 
 *     ImRoseTree.withNodes(1, ImRoseTree.leaf(2), ImRoseTree.leaf(3))
 * 
 * @see ImRoseTreeZipper
 * @see ImRoseTreeIterator
 * @see ImRoseTreeZipperIterator
 * 
 */
public class ImRoseTree<T> implements Iterable<T>
{
    // The element at this node
    final private T element;

    // The children of this element
    final private ImShelf<ImRoseTree<T>> children;

    final private int size;

    private ImRoseTree(T element, ImShelf<ImRoseTree<T>> children)
    {
        this.element = element;
        this.children = children;

        int tempSize = 1;
        for (ImRoseTree<T> t : children)
        {
            tempSize += t.size();
        }
        size = tempSize;
    }

    private ImRoseTree()
    {
        this.element = null;
        this.children = ImShelf.empty();
        size = 0;
    }

    /**
     * A rose tree with element `element` and sub-trees that are `nodes`.
     * 
     * This method is similar to {@link #withNodes(Object, ImRoseTree...)}
     */
    public static <A> ImRoseTree<A> withNodes(A element, Collection<ImRoseTree<A>> nodes)
    {
        return new ImRoseTree<A>(element, ImShelf.onAll(nodes));
    }

    /**
     * A rose tree with element `element` and sub-trees that are `nodes`.
     * 
     * Examples:
     * 
     *     withNodes(1, withElements(2, 3, 4, 5), leaf(6), leaf(7))
     *     
     * produces the tree
     * 
     *     1        
     *     .........
     *     2     6 7
     *     .....    
     *     3 4 5    
     * 
     */
    public static <A> ImRoseTree<A> withNodes(A element, final ImRoseTree<A>... nodes)
    {
        return new ImRoseTree<A>(element, ImShelf.<ImRoseTree<A>> onArray(nodes));
    }

    /**
     * A rose tree with element `element` and sub-trees that are all leaves, with elements `elements`.
     * 
     * This method is similar to {@link #withElements(Object, Object...)}
     */
    public static <A> ImRoseTree<A> withElements(A element, final Collection<A> elements)
    {
        ImShelf<ImRoseTree<A>> shelf = ImShelf.empty();

        for (A el : elements)
        {
            shelf = shelf.adding(leaf(el));
        }

        return new ImRoseTree<A>(element, ImShelf.onAll(shelf));
    }

    /**
     * A rose tree with element `element` and sub-trees that are all leaves, with elements `elements`.
     * 
     * Examples:
     * 
     *     withElements("c", "e", "f", "g");
     *      
     * produces the tree:
     * 
     *     c    
     *     .....
     *     e f g
     *    
     */
    public static <A> ImRoseTree<A> withElements(A element, final A... elements)
    {
        ImShelf<ImRoseTree<A>> shelf = ImShelf.empty();

        for (A el : elements)
        {
            shelf = shelf.adding(leaf(el));
        }

        return new ImRoseTree<A>(element, ImShelf.onAll(shelf));
    }

    /**
     * A rose tree with element `element` and no sub-trees.
     */
    public static <A> ImRoseTree<A> leaf(A element)
    {
        return withNodes(element, Collections.<ImRoseTree<A>> emptyList());
    }

    /**
     * The number of elements/nodes in `this`.
     */
    public int size()
    {
        return size;
    }

    /**
     * `true` if `this` contains `elementToLookFor`.
     * 
     * More formally, returns `true` if and only if `this.toArray()` contains
     * at least one element `e` such that:
     * 
     *     e.equals(elementToLookFor)
     *
     * ### Examples
     * 
     * This code:
     * 
     *     ImRoseTree<String> t = withNodes("a", leaf("b"), withElements("c", withElements("d", "h")));
     *     
     * produces this tree:
     * 
     *         a    
     *         ...  
     *         b c  
     *           ...
     *           d h
     * and then:
     *    
     *     t.contains("d")  =>  true
     *
     */
    public boolean contains(Object elementToLookFor)
    {
        for (Object element : this)
        {
            if (elementToLookFor.equals(element))
                return true;
        }

        return false;
    }

    /**
     * An iterator over the elements in `this` in *pre-order* sequence.
     * 
     * The *pre-order* sequence of a tree is
     * the root element followed by the pre-order sequence of the first child tree, followed by the
     * pre-order sequence of the second child tree and so on.
     *
     * Note that the iterator returned will throw an
     * `UnsupportedOperationException` in response to `remove()`.
     * 
     * The pre-order sequence of this tree:
     * 
     *     a        
     *     .........
     *     b c     d
     *       ..... .
     *       e f g h
     *     
     * is
     * 
     *     a b c e f g d h
     *
     */
    public Iterator<T> iterator()
    {
        return new ImRoseTreeIterator<T>(this);
    }

    /**
     * A string representation of this rose tree.
     * 
     * The root element is shown first. If there are any child trees then it is followed by
     *  `(` then the "toString" of each of the child trees in turn
     * followed by `)`.
     * 
     * ### Examples
     * 
     *     withNodes(1, withElements(2, 3, 4, 5), leaf(6), leaf(7)).toString()  => "1 (2 (3 4 5) 6 7)"
     */
    @Override
    public String toString()
    {
        StringBuilder sb = new StringBuilder();

        sb.append(element);

        if (children.size() > 0)
        {
            sb.append(" (");
            sb.append(TextUtils.join(children, " "));
            sb.append(")");

        }
        return sb.toString();
    }

    /**
     * The element at this node.
     */
    public T getElement()
    {
        return element;
    }

    /**
     * The list (shelf actually) of sub-trees of this tree node.
     * 
     * Note that these are sub-trees of type `ImRoseTree<T>` not elements of type `T`.
     * 
     * ### Examples
     * 
     *     ImRoseTree<String> t = withElements("c", "e", "f", "g");
     *     
     *     t.getSubTrees().size()                       =>  3
     *     t.getSubTrees().get(1).getClass().getName()  =>  "immutablecollections.ImRoseTree"
     * 
     * Note that, if you want to mutate this shelf, you might be tempted to get a zipper on it like this:
     * 
     *     z = t.getSubTrees().getZipper()    // Don't use this to try mutate the sub-trees
     *     
     * and then mutate it. This won't work - because, although you will be able to generate a new
     * mutated shelf of sub-trees, the parent node won't know about it.
     * 
     * To mutate a tree you can use a {@link ImRoseTreeZipper}
     * 
     */
    public ImShelf<ImRoseTree<T>> getSubTrees()
    {
        return children;
    }

    /**
     * A *zipper* on this rose tree.
     */
    public ImRoseTreeZipper<T> getZipper()
    {
        return new ImRoseTreeZipper<T>(this);
    }

    /**
     * An ascii-art diagram of the tree.
     * 
     * ### Examples
     * 
     *     System.err.println(withNodes(1, withElements(2, 3, 4, 5), leaf(6), leaf(7)).toBoxString());
     * 
     * gives:
     * 
     *     1        
     *     .........
     *     2     6 7
     *     .....    
     *     3 4 5  
     * 
     * 
     */
    public String toBoxString()
    {
        return toBox().toString();
    }

    /**
     * The node at index `indexStartingAtOne` in the pre-order sequence.
     * 
     * See {@link #iterator()} for a description of the pre-order sequence.
     * 
     * ### Examples:
     * 
     * Given this code:
     * 
     *     ImRoseTree<String> t = withNodes("a", leaf("b"), withElements("c", "e", "f", "g"), withElements("d", "h"));
     * 
     * which constructs this tree:
     * 
     *     a        
     *     .........
     *     b c     d
     *       ..... .
     *       e f g h
     * 
     * then
     * 
     *      t.getNodeAtIndex(5).getElement()  => "f"
     * 
     */
    public ImRoseTree<T> getNodeAtIndex(final int indexStartingAtOne)
    {
        ImIndexOutOfBoundsException.check(indexStartingAtOne, size, "indexStartingAtOne");

        if (indexStartingAtOne == 1)
        {
            return this;
        }

        int count = 1;
        for (ImRoseTree<T> node : children)
        {
            if (indexStartingAtOne <= count + node.size())
            {
                return node.getNodeAtIndex(indexStartingAtOne - count);
            }

            count += node.size();
        }

        throw new IllegalStateException("Can't get here");
    }

    protected TextBox toBox()
    {
        TextBox myTopBox = LeafTextBox.with(element.toString());

        if (getSubTrees().isEmpty())
            return myTopBox;

        LinkedList<TextBox> boxes = new LinkedList<TextBox>();

        for (ImRoseTree<T> t : getSubTrees())
        {
            boxes.add(t.toBox());
        }

        TextBox space = LeafTextBox.with(" ");

        // Intersperse the child boxes with spaces
        LeftRightBox childBoxes = LeftRightBox.withAll(ImList.onAll(boxes).intersperse(space));

        /**
         * Calculate the line of dots underneath the top box.
         * 
         * If we make the width of the dots the width of all my child boxes it will look like this: (eg)
         *    1    
         *    .....
         *    2 3  
         *      ...
         *      4 5
         *   
         * We reduce the width to only extend as far as the end of the top box of the last child. It will
         * now look like this:
         * 
         *    1    
         *    ...  
         *    2 3  
         *      ...
         *      4 5
         * 
         * 
         */
        int rawWidth = childBoxes.getWidth();

        ImRoseTree<T> lastChild = getSubTrees().get1(getSubTrees().size());

        int reducedWidth = rawWidth - (boxes.getLast().getWidth() - lastChild.element.toString().length());

        final TextBox dots = LeafTextBox.with(TextUtils.repeat(".", reducedWidth));

        return TopDownBox.with(myTopBox, dots, childBoxes);
    }

    /**
     * The rose tree that has the same sub-trees as `this` but has `newElement` as its element.
     * 
     * If `newElement == element` then return `this`
     * 
     */
    public ImRoseTree<T> replaceElement(T newElement)
    {
        return newElement == element
                ? this
                : withNodes(newElement, children);
    }

    /**
     * An iterator on a zipper on this rose tree.
     */
    public Iterator<ImRoseTreeZipper<T>> getZipperIterator()
    {
        return getZipper().iterator();
    }

    /**
     * 
     * The rose tree that has the same shape as this and where the i-th element is the
     * result of evaluating the single argument function `fn` on
     * the i-th element of `this`.
     * 
     * ### Examples
     * 
     *     ImRoseTree<String> t = withElements("c", "e", "f", "g");    
     *     
     *     Function1<String> toUpperFn = FnFactory.on(String.class).getFn(String.class, "toUpperCase");
     *     
     *     t.map(toUpperFn)  =>  C (E F G)
     *     
     */
    public <O> ImRoseTree<O> map(Function1<O> fn)
    {
        return new ImRoseTree<O>(fn.invoke(element), children.map(lift(fn)));
    }

    private <O> Function1<ImRoseTree<O>> lift(Function1<O> fn)
    {
        return (Function1<ImRoseTree<O>>) mapFn.flip().invoke(fn);
    }

    // The instance function map
    private static Function2<?> mapFn = FnFactory.on(ImRoseTree.class).getFn(ImRoseTree.class, "map", Function1.class);

}
